{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/NielsRogge/Transformers-Tutorials/blob/master/Custom_Named_Entity_Recognition_with_BERT.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "MyETdB-dkBsX"
   },
   "source": [
    "## **Fine-tuning BERT for named-entity recognition**\n",
    "\n",
    "In this notebook, we are going to use **BertForTokenClassification** which is included in the [Transformers library](https://github.com/huggingface/transformers) by HuggingFace. This model has BERT as its base architecture, with a token classification head on top, allowing it to make predictions at the token level, rather than the sequence level. Named entity recognition is typically treated as a token classification problem, so that's what we are going to use it for.\n",
    "\n",
    "This tutorial uses the idea of **transfer learning**, i.e. first pretraining a large neural network in an unsupervised way, and then fine-tuning that neural network on a task of interest. In this case, BERT is a neural network pretrained on 2 tasks: masked language modeling and next sentence prediction. Now, we are going to fine-tune this network on a NER dataset. Fine-tuning is supervised learning, so this means we will need a labeled dataset.\n",
    "\n",
    "If you want to know more about BERT, I suggest the following resources:\n",
    "* the original [paper](https://arxiv.org/abs/1810.04805)\n",
    "* Jay Allamar's [blog post](http://jalammar.github.io/illustrated-bert/) as well as his [tutorial](http://jalammar.github.io/a-visual-guide-to-using-bert-for-the-first-time/)\n",
    "* Chris Mccormick's [Youtube channel](https://www.youtube.com/channel/UCoRX98PLOsaN8PtekB9kWrw)\n",
    "* Abbishek Kumar Mishra's [Youtube channel](https://www.youtube.com/user/abhisheksvnit)\n",
    "\n",
    "The following notebook largely follows the same structure as the tutorials by Abhishek Kumar Mishra. For his tutorials on the Transformers library, see his [Github repository](https://github.com/abhimishra91/transformers-tutorials).\n",
    "\n",
    "NOTE: this notebook assumes basic knowledge about deep learning, BERT, and native PyTorch. If you want to learn more Python, deep learning and PyTorch, I highly recommend cs231n by Stanford University and the FastAI course by Jeremy Howard et al. Both are freely available on the web.  \n",
    "\n",
    "Now, let's move on to the real stuff!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "e7wfLWyYkvDi"
   },
   "source": [
    "#### **Importing Python Libraries and preparing the environment**\n",
    "\n",
    "This notebook assumes that you have the following libraries installed:\n",
    "* pandas\n",
    "* numpy\n",
    "* sklearn\n",
    "* pytorch\n",
    "* transformers\n",
    "* seqeval\n",
    "\n",
    "As we are running this in Google Colab, the only libraries we need to additionally install are transformers and seqeval (GPU version):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "#!pip install transformers seqeval[gpu]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "IEnlUbgm8z3B"
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "from sklearn.metrics import accuracy_score\n",
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from transformers import BertTokenizer, BertConfig, BertForTokenClassification"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Jzq1w3L1K5M-"
   },
   "source": [
    "As deep learning can be accellerated a lot using a GPU instead of a CPU, make sure you can run this notebook in a GPU runtime (which Google Colab provides for free! - check \"Runtime\" - \"Change runtime type\" - and set the hardware accelerator to \"GPU\").\n",
    "\n",
    "We can set the default device to GPU using the following code (if it prints \"cuda\", it means the GPU has been recognized):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "Sm1krxJtKxpx",
    "outputId": "c0cf11ba-17ff-4a35-da11-5b7b0c4cb11d",
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda\n"
     ]
    }
   ],
   "source": [
    "from torch import cuda\n",
    "device = 'cuda' if cuda.is_available() else 'cpu'\n",
    "print(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ahwMsmyG5ZPE"
   },
   "source": [
    "#### **Downloading and preprocessing the data**\n",
    "Named entity recognition (NER) uses a specific annotation scheme, which is defined (at least for European languages) at the *word* level. An annotation scheme that is widely used is called **[IOB-tagging](https://en.wikipedia.org/wiki/Inside%E2%80%93outside%E2%80%93beginning_(tagging))**, which stands for Inside-Outside-Beginning. Each tag indicates whether the corresponding word is *inside*, *outside* or at the *beginning* of a specific named entity. The reason this is used is because named entities usually comprise more than 1 word. \n",
    "\n",
    "Let's have a look at an example. If you have a sentence like \"Barack Obama was born in Hawaï\", then the corresponding tags would be   [B-PERS, I-PERS, O, O, O, B-GEO]. B-PERS means that the word \"Barack\" is the beginning of a person, I-PERS means that the word \"Obama\" is inside a person, \"O\" means that the word \"was\" is outside a named entity, and so on. So one typically has as many tags as there are words in a sentence.\n",
    "\n",
    "So if you want to train a deep learning model for NER, it requires that you have your data in this IOB format (or similar formats such as [BILOU](https://stackoverflow.com/questions/17116446/what-do-the-bilou-tags-mean-in-named-entity-recognition)). There exist many annotation tools which let you create these kind of annotations automatically (such as Spacy's [Prodigy](https://prodi.gy/), [Tagtog](https://docs.tagtog.net/) or [Doccano](https://github.com/doccano/doccano)). You can also use Spacy's [biluo_tags_from_offsets](https://spacy.io/api/goldparse#biluo_tags_from_offsets) function to convert annotations at the character level to IOB format.\n",
    "\n",
    "Here, we will use a NER dataset from [Kaggle](https://www.kaggle.com/namanj27/ner-dataset) that is already in IOB format. One has to go to this web page, download the dataset, unzip it, and upload the csv file to this notebook. Let's print out the first few rows of this csv file:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 195
    },
    "colab_type": "code",
    "id": "deLB9HVX5I6F",
    "outputId": "57531522-0650-4a3c-dc6a-b48a9348a17a"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Sentence #</th>\n",
       "      <th>Word</th>\n",
       "      <th>POS</th>\n",
       "      <th>Tag</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Sentence: 1</td>\n",
       "      <td>Thousands</td>\n",
       "      <td>NNS</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>NaN</td>\n",
       "      <td>of</td>\n",
       "      <td>IN</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>NaN</td>\n",
       "      <td>demonstrators</td>\n",
       "      <td>NNS</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>NaN</td>\n",
       "      <td>have</td>\n",
       "      <td>VBP</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>NaN</td>\n",
       "      <td>marched</td>\n",
       "      <td>VBN</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    Sentence #           Word  POS Tag\n",
       "0  Sentence: 1      Thousands  NNS   O\n",
       "1          NaN             of   IN   O\n",
       "2          NaN  demonstrators  NNS   O\n",
       "3          NaN           have  VBP   O\n",
       "4          NaN        marched  VBN   O"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = pd.read_csv(\"ner_datasetreference.csv\", encoding='unicode_escape')\n",
    "data.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ucYjhq6uRAmY"
   },
   "source": [
    "Let's check how many sentences and words (and corresponding tags) there are in this dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 101
    },
    "colab_type": "code",
    "id": "6gMibEJXTKDw",
    "outputId": "a7039b75-8800-4a1e-8554-89bd41168683"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sentence #      47959\n",
       "Word          1048575\n",
       "POS           1048575\n",
       "Tag           1048575\n",
       "dtype: int64"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.count()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see, there are approximately 48,000 sentences in the dataset, comprising more than 1 million words and tags (quite huge!). This corresponds to approximately 20 words per sentence. \n",
    "\n",
    "Let's have a look at the different NER tags, and their frequency: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of tags: 17\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "O        887908\n",
       "B-geo     37644\n",
       "B-tim     20333\n",
       "B-org     20143\n",
       "I-per     17251\n",
       "B-per     16990\n",
       "I-org     16784\n",
       "B-gpe     15870\n",
       "I-geo      7414\n",
       "I-tim      6528\n",
       "B-art       402\n",
       "B-eve       308\n",
       "I-art       297\n",
       "I-eve       253\n",
       "B-nat       201\n",
       "I-gpe       198\n",
       "I-nat        51\n",
       "Name: Tag, dtype: int64"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(\"Number of tags: {}\".format(len(data.Tag.unique())))\n",
    "frequencies = data.Tag.value_counts()\n",
    "frequencies"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are 8 category tags, each with a \"beginning\" and \"inside\" variant, and the \"outside\" tag. It is not really clear what these tags mean - \"geo\" probably stands for geographical entity, \"gpe\" for geopolitical entity, and so on. They do not seem to correspond with what the publisher says on Kaggle. Some tags seem to be underrepresented. Let's print them by frequency (highest to lowest): "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[('geo', 45058), ('org', 36927), ('per', 34241), ('tim', 26861), ('gpe', 16068), ('art', 699), ('eve', 561), ('nat', 252)]\n"
     ]
    }
   ],
   "source": [
    "tags = {}\n",
    "for tag, count in zip(frequencies.index, frequencies):\n",
    "    if tag != \"O\":\n",
    "        if tag[2:5] not in tags.keys():\n",
    "            tags[tag[2:5]] = count\n",
    "        else:\n",
    "            tags[tag[2:5]] += count\n",
    "    continue\n",
    "\n",
    "print(sorted(tags.items(), key=lambda x: x[1], reverse=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's remove \"art\", \"eve\" and \"nat\" named entities, as performance on them will probably be not comparable to the other named entities. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Sentence #</th>\n",
       "      <th>Word</th>\n",
       "      <th>POS</th>\n",
       "      <th>Tag</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Sentence: 1</td>\n",
       "      <td>Thousands</td>\n",
       "      <td>NNS</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>NaN</td>\n",
       "      <td>of</td>\n",
       "      <td>IN</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>NaN</td>\n",
       "      <td>demonstrators</td>\n",
       "      <td>NNS</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>NaN</td>\n",
       "      <td>have</td>\n",
       "      <td>VBP</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>NaN</td>\n",
       "      <td>marched</td>\n",
       "      <td>VBN</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    Sentence #           Word  POS Tag\n",
       "0  Sentence: 1      Thousands  NNS   O\n",
       "1          NaN             of   IN   O\n",
       "2          NaN  demonstrators  NNS   O\n",
       "3          NaN           have  VBP   O\n",
       "4          NaN        marched  VBN   O"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "entities_to_remove = [\"B-art\", \"I-art\", \"B-eve\", \"I-eve\", \"B-nat\", \"I-nat\"]\n",
    "data = data[~data.Tag.isin(entities_to_remove)]\n",
    "data.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "mskU4h0oRKEF"
   },
   "source": [
    "Now, we have to ask ourself the question: what is a training example in the case of NER, which is provided in a single forward pass? A training example is typically a **sentence**, with corresponding IOB tags. Let's group the words and corresponding tags by sentence:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 195
    },
    "colab_type": "code",
    "id": "zkW2vNcO-uMH",
    "outputId": "855c09f8-1345-4851-a060-1f3a8290d4fe"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Sentence #</th>\n",
       "      <th>Word</th>\n",
       "      <th>POS</th>\n",
       "      <th>Tag</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Sentence: 1</td>\n",
       "      <td>Thousands</td>\n",
       "      <td>NNS</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Sentence: 1</td>\n",
       "      <td>of</td>\n",
       "      <td>IN</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Sentence: 1</td>\n",
       "      <td>demonstrators</td>\n",
       "      <td>NNS</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Sentence: 1</td>\n",
       "      <td>have</td>\n",
       "      <td>VBP</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Sentence: 1</td>\n",
       "      <td>marched</td>\n",
       "      <td>VBN</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    Sentence #           Word  POS Tag\n",
       "0  Sentence: 1      Thousands  NNS   O\n",
       "1  Sentence: 1             of   IN   O\n",
       "2  Sentence: 1  demonstrators  NNS   O\n",
       "3  Sentence: 1           have  VBP   O\n",
       "4  Sentence: 1        marched  VBN   O"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# pandas has a very handy \"forward fill\" function to fill missing values based on the last upper non-nan value\n",
    "data = data.fillna(method='ffill')\n",
    "data.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 195
    },
    "colab_type": "code",
    "id": "Hmd-ow389k6Y",
    "outputId": "51e3a794-9847-4462-fd79-e280c00d81ce"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Sentence #</th>\n",
       "      <th>Word</th>\n",
       "      <th>POS</th>\n",
       "      <th>Tag</th>\n",
       "      <th>sentence</th>\n",
       "      <th>word_labels</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Sentence: 1</td>\n",
       "      <td>Thousands</td>\n",
       "      <td>NNS</td>\n",
       "      <td>O</td>\n",
       "      <td>Thousands of demonstrators have marched throug...</td>\n",
       "      <td>O,O,O,O,O,O,B-geo,O,O,O,O,O,B-geo,O,O,O,O,O,B-...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Sentence: 1</td>\n",
       "      <td>of</td>\n",
       "      <td>IN</td>\n",
       "      <td>O</td>\n",
       "      <td>Thousands of demonstrators have marched throug...</td>\n",
       "      <td>O,O,O,O,O,O,B-geo,O,O,O,O,O,B-geo,O,O,O,O,O,B-...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Sentence: 1</td>\n",
       "      <td>demonstrators</td>\n",
       "      <td>NNS</td>\n",
       "      <td>O</td>\n",
       "      <td>Thousands of demonstrators have marched throug...</td>\n",
       "      <td>O,O,O,O,O,O,B-geo,O,O,O,O,O,B-geo,O,O,O,O,O,B-...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Sentence: 1</td>\n",
       "      <td>have</td>\n",
       "      <td>VBP</td>\n",
       "      <td>O</td>\n",
       "      <td>Thousands of demonstrators have marched throug...</td>\n",
       "      <td>O,O,O,O,O,O,B-geo,O,O,O,O,O,B-geo,O,O,O,O,O,B-...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Sentence: 1</td>\n",
       "      <td>marched</td>\n",
       "      <td>VBN</td>\n",
       "      <td>O</td>\n",
       "      <td>Thousands of demonstrators have marched throug...</td>\n",
       "      <td>O,O,O,O,O,O,B-geo,O,O,O,O,O,B-geo,O,O,O,O,O,B-...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    Sentence #           Word  POS Tag  \\\n",
       "0  Sentence: 1      Thousands  NNS   O   \n",
       "1  Sentence: 1             of   IN   O   \n",
       "2  Sentence: 1  demonstrators  NNS   O   \n",
       "3  Sentence: 1           have  VBP   O   \n",
       "4  Sentence: 1        marched  VBN   O   \n",
       "\n",
       "                                            sentence  \\\n",
       "0  Thousands of demonstrators have marched throug...   \n",
       "1  Thousands of demonstrators have marched throug...   \n",
       "2  Thousands of demonstrators have marched throug...   \n",
       "3  Thousands of demonstrators have marched throug...   \n",
       "4  Thousands of demonstrators have marched throug...   \n",
       "\n",
       "                                         word_labels  \n",
       "0  O,O,O,O,O,O,B-geo,O,O,O,O,O,B-geo,O,O,O,O,O,B-...  \n",
       "1  O,O,O,O,O,O,B-geo,O,O,O,O,O,B-geo,O,O,O,O,O,B-...  \n",
       "2  O,O,O,O,O,O,B-geo,O,O,O,O,O,B-geo,O,O,O,O,O,B-...  \n",
       "3  O,O,O,O,O,O,B-geo,O,O,O,O,O,B-geo,O,O,O,O,O,B-...  \n",
       "4  O,O,O,O,O,O,B-geo,O,O,O,O,O,B-geo,O,O,O,O,O,B-...  "
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# let's create a new column called \"sentence\" which groups the words by sentence \n",
    "data['sentence'] = data[['Sentence #','Word','Tag']].groupby(['Sentence #'])['Word'].transform(lambda x: ' '.join(x))\n",
    "# let's also create a new column called \"word_labels\" which groups the tags by sentence \n",
    "data['word_labels'] = data[['Sentence #','Word','Tag']].groupby(['Sentence #'])['Tag'].transform(lambda x: ','.join(x))\n",
    "data.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "JsjhdQbE-Lve"
   },
   "source": [
    "Let's have a look at the different NER tags. \n",
    "\n",
    "We create 2 dictionaries: one that maps individual tags to indices, and one that maps indices to their individual tags. This is necessary in order to create the labels (as computers work with numbers = indices, rather than words = tags) - see further in this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 302
    },
    "colab_type": "code",
    "id": "CFRDM8WsQXvL",
    "outputId": "1e9d7dcb-f93c-4d71-e342-557c1709206b"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'O': 0,\n",
       " 'B-geo': 1,\n",
       " 'B-gpe': 2,\n",
       " 'B-per': 3,\n",
       " 'I-geo': 4,\n",
       " 'B-org': 5,\n",
       " 'I-org': 6,\n",
       " 'B-tim': 7,\n",
       " 'I-per': 8,\n",
       " 'I-gpe': 9,\n",
       " 'I-tim': 10}"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "labels_to_ids = {k: v for v, k in enumerate(data.Tag.unique())}\n",
    "ids_to_labels = {v: k for v, k in enumerate(data.Tag.unique())}\n",
    "labels_to_ids"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "J08Cvk_USgbM"
   },
   "source": [
    "As we can see, there are now only 10 different tags. \n",
    "\n",
    "Let's only keep the \"sentence\" and \"word_labels\" columns, and drop duplicates:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 195
    },
    "colab_type": "code",
    "id": "SrEgd4PZUgmF",
    "outputId": "9411df51-39d8-4757-a98d-acdd3319aa08"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>sentence</th>\n",
       "      <th>word_labels</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Thousands of demonstrators have marched throug...</td>\n",
       "      <td>O,O,O,O,O,O,B-geo,O,O,O,O,O,B-geo,O,O,O,O,O,B-...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Families of soldiers killed in the conflict jo...</td>\n",
       "      <td>O,O,O,O,O,O,O,O,O,O,O,O,O,O,O,O,O,O,B-per,O,O,...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>They marched from the Houses of Parliament to ...</td>\n",
       "      <td>O,O,O,O,O,O,O,O,O,O,O,B-geo,I-geo,O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Police put the number of marchers at 10,000 wh...</td>\n",
       "      <td>O,O,O,O,O,O,O,O,O,O,O,O,O,O,O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>The protest comes on the eve of the annual con...</td>\n",
       "      <td>O,O,O,O,O,O,O,O,O,O,O,B-geo,O,O,B-org,I-org,O,...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                            sentence  \\\n",
       "0  Thousands of demonstrators have marched throug...   \n",
       "1  Families of soldiers killed in the conflict jo...   \n",
       "2  They marched from the Houses of Parliament to ...   \n",
       "3  Police put the number of marchers at 10,000 wh...   \n",
       "4  The protest comes on the eve of the annual con...   \n",
       "\n",
       "                                         word_labels  \n",
       "0  O,O,O,O,O,O,B-geo,O,O,O,O,O,B-geo,O,O,O,O,O,B-...  \n",
       "1  O,O,O,O,O,O,O,O,O,O,O,O,O,O,O,O,O,O,B-per,O,O,...  \n",
       "2                O,O,O,O,O,O,O,O,O,O,O,B-geo,I-geo,O  \n",
       "3                      O,O,O,O,O,O,O,O,O,O,O,O,O,O,O  \n",
       "4  O,O,O,O,O,O,O,O,O,O,O,B-geo,O,O,B-org,I-org,O,...  "
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = data[[\"sentence\", \"word_labels\"]].drop_duplicates().reset_index(drop=True)\n",
    "data.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "r3ArUiVRqw0C",
    "outputId": "b78ec91c-88a4-4433-9bd4-c9f726dc53e3"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "47571"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "U8obZumRTBrT"
   },
   "source": [
    "Let's verify that a random sentence and its corresponding tags are correct:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "eUvupomW_fbe",
    "outputId": "926daac8-b2a0-436d-f98c-cbe93d1b00ed"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Bedfordshire police said Tuesday that Omar Khayam was arrested in Bedford for breaching the conditions of his parole .'"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.iloc[41].sentence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "0dLyY3Oi_lvp",
    "outputId": "5d61d25a-6d01-4249-ba53-30d268f5a330"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'B-gpe,O,O,B-tim,O,B-per,I-per,O,O,O,B-geo,O,O,O,O,O,O,O,O'"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.iloc[41].word_labels"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "f5EHpuB78pIa"
   },
   "source": [
    "#### **Preparing the dataset and dataloader**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "15x7zmZnTgFx"
   },
   "source": [
    "Now that our data is preprocessed, we can turn it into PyTorch tensors such that we can provide it to the model. Let's start by defining some key variables that will be used later on in the training/evaluation process:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "lgNSM8Xz79Mg",
    "tags": []
   },
   "outputs": [],
   "source": [
    "MAX_LEN = 128\n",
    "TRAIN_BATCH_SIZE = 4\n",
    "VALID_BATCH_SIZE = 2\n",
    "EPOCHS = 1\n",
    "LEARNING_RATE = 1e-05\n",
    "MAX_GRAD_NORM = 10\n",
    "tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "wPYV2Ld6Tr5I"
   },
   "source": [
    "A tricky part of NER with BERT is that BERT relies on **wordpiece tokenization**, rather than word tokenization. This means that we should also define the labels at the wordpiece-level, rather than the word-level! \n",
    "\n",
    "For example, if you have word like \"Washington\" which is labeled as \"b-gpe\", but it gets tokenized to \"Wash\", \"##ing\", \"##ton\", then we will have to propagate the word’s original label to all of its wordpieces: \"b-gpe\", \"b-gpe\", \"b-gpe\". The model should be able to produce the correct labels for each individual wordpiece. The function below (taken from [here](https://github.com/chambliss/Multilingual_NER/blob/master/python/utils/main_utils.py#L118)) implements this.\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "RNzSgZTfGUd8"
   },
   "outputs": [],
   "source": [
    "def tokenize_and_preserve_labels(sentence, text_labels, tokenizer):\n",
    "    \"\"\"\n",
    "    Word piece tokenization makes it difficult to match word labels\n",
    "    back up with individual word pieces. This function tokenizes each\n",
    "    word one at a time so that it is easier to preserve the correct\n",
    "    label for each subword. It is, of course, a bit slower in processing\n",
    "    time, but it will help our model achieve higher accuracy.\n",
    "    \"\"\"\n",
    "\n",
    "    tokenized_sentence = []\n",
    "    labels = []\n",
    "\n",
    "    sentence = sentence.strip()\n",
    "\n",
    "    for word, label in zip(sentence.split(), text_labels.split(\",\")):\n",
    "\n",
    "        # Tokenize the word and count # of subwords the word is broken into\n",
    "        tokenized_word = tokenizer.tokenize(word)\n",
    "        n_subwords = len(tokenized_word)\n",
    "\n",
    "        # Add the tokenized word to the final tokenized word list\n",
    "        tokenized_sentence.extend(tokenized_word)\n",
    "\n",
    "        # Add the same label to the new list of labels `n_subwords` times\n",
    "        labels.extend([label] * n_subwords)\n",
    "\n",
    "    return tokenized_sentence, labels"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ez7qlFHl56ZW"
   },
   "source": [
    "Note that this is a **design decision**. You could also decide to only label the first wordpiece of each word and let the model only learn this (this is what was done in the original BERT paper, see Github discussion [here](https://github.com/huggingface/transformers/issues/64#issuecomment-443703063)). Another design decision could be to give the first wordpiece of each word the original word label, and then use the label “X” for all subsequent subwords of that word.\n",
    "\n",
    "All of them lead to good performance.\n",
    "\n",
    "Next, we define a regular PyTorch [dataset class](https://pytorch.org/docs/stable/data.html) (which transforms examples of a dataframe to PyTorch tensors). Here, each sentence gets tokenized, the special tokens that BERT expects are added, the tokens are padded or truncated based on the max length of the model, the attention mask is created and the labels are created based on the dictionary which we defined above. \n",
    "\n",
    "For more information about BERT's inputs, see [here](https://huggingface.co/transformers/glossary.html).  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "aJty_Abw8_xK"
   },
   "outputs": [],
   "source": [
    "class dataset(Dataset):\n",
    "    def __init__(self, dataframe, tokenizer, max_len):\n",
    "        self.len = len(dataframe)\n",
    "        self.data = dataframe\n",
    "        self.tokenizer = tokenizer\n",
    "        self.max_len = max_len\n",
    "        \n",
    "    def __getitem__(self, index):\n",
    "        # step 1: tokenize (and adapt corresponding labels)\n",
    "        sentence = self.data.sentence[index]  \n",
    "        word_labels = self.data.word_labels[index]  \n",
    "        tokenized_sentence, labels = tokenize_and_preserve_labels(sentence, word_labels, self.tokenizer)\n",
    "        \n",
    "        # step 2: add special tokens (and corresponding labels)\n",
    "        tokenized_sentence = [\"[CLS]\"] + tokenized_sentence + [\"[SEP]\"] # add special tokens\n",
    "        labels.insert(0, \"O\") # add outside label for [CLS] token\n",
    "        labels.insert(-1, \"O\") # add outside label for [SEP] token\n",
    "\n",
    "        # step 3: truncating/padding\n",
    "        maxlen = self.max_len\n",
    "\n",
    "        if (len(tokenized_sentence) > maxlen):\n",
    "          # truncate\n",
    "          tokenized_sentence = tokenized_sentence[:maxlen]\n",
    "          labels = labels[:maxlen]\n",
    "        else:\n",
    "          # pad\n",
    "          tokenized_sentence = tokenized_sentence + ['[PAD]'for _ in range(maxlen - len(tokenized_sentence))]\n",
    "          labels = labels + [\"O\" for _ in range(maxlen - len(labels))]\n",
    "\n",
    "        # step 4: obtain the attention mask\n",
    "        attn_mask = [1 if tok != '[PAD]' else 0 for tok in tokenized_sentence]\n",
    "        \n",
    "        # step 5: convert tokens to input ids\n",
    "        ids = self.tokenizer.convert_tokens_to_ids(tokenized_sentence)\n",
    "\n",
    "        label_ids = [labels_to_ids[label] for label in labels]\n",
    "        # the following line is deprecated\n",
    "        #label_ids = [label if label != 0 else -100 for label in label_ids]\n",
    "        \n",
    "        return {\n",
    "              'ids': torch.tensor(ids, dtype=torch.long),\n",
    "              'mask': torch.tensor(attn_mask, dtype=torch.long),\n",
    "              #'token_type_ids': torch.tensor(token_ids, dtype=torch.long),\n",
    "              'targets': torch.tensor(label_ids, dtype=torch.long)\n",
    "        } \n",
    "    \n",
    "    def __len__(self):\n",
    "        return self.len"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "hTP7zuWGWGUd"
   },
   "source": [
    "Now, based on the class we defined above, we can create 2 datasets, one for training and one for testing. Let's use a 80/20 split:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 67
    },
    "colab_type": "code",
    "id": "jrkdZBLYHVcB",
    "outputId": "160536bb-9659-490e-db38-c5c7ff66b351",
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FULL Dataset: (47571, 2)\n",
      "TRAIN Dataset: (38057, 2)\n",
      "TEST Dataset: (9514, 2)\n"
     ]
    }
   ],
   "source": [
    "train_size = 0.8\n",
    "train_dataset = data.sample(frac=train_size,random_state=200)\n",
    "test_dataset = data.drop(train_dataset.index).reset_index(drop=True)\n",
    "train_dataset = train_dataset.reset_index(drop=True)\n",
    "\n",
    "print(\"FULL Dataset: {}\".format(data.shape))\n",
    "print(\"TRAIN Dataset: {}\".format(train_dataset.shape))\n",
    "print(\"TEST Dataset: {}\".format(test_dataset.shape))\n",
    "\n",
    "training_set = dataset(train_dataset, tokenizer, MAX_LEN)\n",
    "testing_set = dataset(test_dataset, tokenizer, MAX_LEN)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Ptv5AT_iTb7W"
   },
   "source": [
    "Let's have a look at the first training example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "colab_type": "code",
    "id": "phmPylgAm8Xy",
    "outputId": "27a7cee8-c920-4602-c539-b117ce6236ec"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'ids': tensor([  101, 23564, 21030,  2099,  4967,  2001,  9388,  1011,  6109,  2005,\n",
       "          2634,  1012,   102,     0,     0,     0,     0,     0,     0,     0,\n",
       "             0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "             0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "             0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "             0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "             0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "             0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "             0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "             0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "             0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "             0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "             0,     0,     0,     0,     0,     0,     0,     0]),\n",
       " 'mask': tensor([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "         0, 0, 0, 0, 0, 0, 0, 0]),\n",
       " 'targets': tensor([0, 3, 3, 3, 8, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "         0, 0, 0, 0, 0, 0, 0, 0])}"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "training_set[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "VvU4nzL2W2Xo"
   },
   "source": [
    "Let's verify that the input ids and corresponding targets are correct:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "colab_type": "code",
    "id": "DWgnNJrYW2GP",
    "outputId": "2b5b5b04-b93c-4354-b104-10b96c7223ff",
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[CLS]       0\n",
      "za          3\n",
      "##hee       3\n",
      "##r         3\n",
      "khan        8\n",
      "was         0\n",
      "mar         0\n",
      "-           0\n",
      "93          0\n",
      "for         0\n",
      "india       1\n",
      ".           0\n",
      "[SEP]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n",
      "[PAD]       0\n"
     ]
    }
   ],
   "source": [
    "for token, label in zip(tokenizer.convert_ids_to_tokens(training_set[0][\"ids\"]), training_set[0][\"targets\"]):\n",
    "  print('{0:10}  {1}'.format(token, label))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Ky68FcTgWnfN"
   },
   "source": [
    "Now, let's define the corresponding PyTorch dataloaders:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "KIw793myWOmi"
   },
   "outputs": [],
   "source": [
    "train_params = {'batch_size': TRAIN_BATCH_SIZE,\n",
    "                'shuffle': True,\n",
    "                'num_workers': 0\n",
    "                }\n",
    "\n",
    "test_params = {'batch_size': VALID_BATCH_SIZE,\n",
    "                'shuffle': True,\n",
    "                'num_workers': 0\n",
    "                }\n",
    "\n",
    "training_loader = DataLoader(training_set, **train_params)\n",
    "testing_loader = DataLoader(testing_set, **test_params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "73OzU7oXRxR8"
   },
   "source": [
    "#### **Defining the model**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "T-iGhnhdLNdP"
   },
   "source": [
    "Here we define the model, BertForTokenClassification, and load it with the pretrained weights of \"bert-base-uncased\". The only thing we need to additionally specify is the number of labels (as this will determine the architecture of the classification head).\n",
    "\n",
    "Note that only the base layers are initialized with the pretrained weights. The token classification head of top has just randomly initialized weights, which we will train, together with the pretrained weights, using our labelled dataset. This is also printed as a warning when you run the code cell below.\n",
    "\n",
    "Then, we move the model to the GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "colab_type": "code",
    "id": "cB9MR3KcWXUs",
    "outputId": "c0d93a4f-cd40-4f6b-edab-6dd51f2ee33b",
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at bert-base-uncased were not used when initializing BertForTokenClassification: ['cls.predictions.bias', 'cls.predictions.transform.dense.weight', 'cls.predictions.transform.dense.bias', 'cls.predictions.decoder.weight', 'cls.seq_relationship.weight', 'cls.seq_relationship.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.transform.LayerNorm.bias']\n",
      "- This IS expected if you are initializing BertForTokenClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPretraining model).\n",
      "- This IS NOT expected if you are initializing BertForTokenClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some weights of BertForTokenClassification were not initialized from the model checkpoint at bert-base-uncased and are newly initialized: ['classifier.weight', 'classifier.bias']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "BertForTokenClassification(\n",
       "  (bert): BertModel(\n",
       "    (embeddings): BertEmbeddings(\n",
       "      (word_embeddings): Embedding(30522, 768, padding_idx=0)\n",
       "      (position_embeddings): Embedding(512, 768)\n",
       "      (token_type_embeddings): Embedding(2, 768)\n",
       "      (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "      (dropout): Dropout(p=0.1, inplace=False)\n",
       "    )\n",
       "    (encoder): BertEncoder(\n",
       "      (layer): ModuleList(\n",
       "        (0): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (1): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (2): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (3): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (4): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (5): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (6): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (7): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (8): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (9): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (10): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (11): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (pooler): BertPooler(\n",
       "      (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "      (activation): Tanh()\n",
       "    )\n",
       "  )\n",
       "  (dropout): Dropout(p=0.1, inplace=False)\n",
       "  (classifier): Linear(in_features=768, out_features=11, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = BertForTokenClassification.from_pretrained('bert-base-uncased', num_labels=len(labels_to_ids))\n",
    "model.to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Pp7Yl4JyWhDj"
   },
   "source": [
    "#### **Training the model**\n",
    "\n",
    "Before training the model, let's perform a sanity check, which I learned thanks to Andrej Karpathy's wonderful [cs231n course](http://cs231n.stanford.edu/) at Stanford (see also his [blog post about debugging neural networks](http://karpathy.github.io/2019/04/25/recipe/)). The initial loss of your model should be close to -ln(1/number of classes) = -ln(1/17) = 2.83. \n",
    "\n",
    "Why? Because we are using cross entropy loss. The cross entropy loss is defined as -ln(probability score of the model for the correct class). In the beginning, the weights are random, so the probability distribution for all of the classes for a given token will be uniform, meaning that the probability for the correct class will be near 1/17. The loss for a given token will thus be -ln(1/17). As PyTorch's [CrossEntropyLoss](https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html) (which is used by `BertForTokenClassification`) uses *mean reduction* by default, it will compute the mean loss for each of the tokens in the sequence (in other words, for all of the 512 tokens). The mean of 512 times -log(1/17) is, you guessed it, -log(1/17).  \n",
    "\n",
    "Let's verify this:\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "eqAN7YVIjKTr",
    "outputId": "0a6acb55-a46e-46b6-e562-e4205a9857cd"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(2.3348, device='cuda:0', grad_fn=<NllLossBackward>)"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ids = training_set[0][\"ids\"].unsqueeze(0)\n",
    "mask = training_set[0][\"mask\"].unsqueeze(0)\n",
    "targets = training_set[0][\"targets\"].unsqueeze(0)\n",
    "ids = ids.to(device)\n",
    "mask = mask.to(device)\n",
    "targets = targets.to(device)\n",
    "outputs = model(input_ids=ids, attention_mask=mask, labels=targets)\n",
    "initial_loss = outputs[0]\n",
    "initial_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "yLdwsru9Mh7U"
   },
   "source": [
    "This looks good. Let's also verify that the logits of the neural network have a shape of (batch_size, sequence_length, num_labels):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "X-z6YCpGnvfj",
    "outputId": "b873271e-b7f8-48b6-f4f2-b63c3ccccbe5"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 128, 11])"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tr_logits = outputs[1]\n",
    "tr_logits.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "kwDLXxOVOCvD"
   },
   "source": [
    "Next, we define the optimizer. Here, we are just going to use Adam with a default learning rate. One can also decide to use more advanced ones such as AdamW (Adam with weight decay fix), which is [included](https://huggingface.co/transformers/main_classes/optimizer_schedules.html) in the Transformers repository, and a learning rate scheduler, but we are not going to do that here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "kznSQfGIWdU4"
   },
   "outputs": [],
   "source": [
    "optimizer = torch.optim.Adam(params=model.parameters(), lr=LEARNING_RATE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "vZQ8JMF0NOe1"
   },
   "source": [
    "Now let's define a regular PyTorch training function. It is partly based on [a really good repository about multilingual NER](https://github.com/chambliss/Multilingual_NER/blob/master/python/utils/main_utils.py#L344)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "GLFivpkwW1HY"
   },
   "outputs": [],
   "source": [
    "# Defining the training function on the 80% of the dataset for tuning the bert model\n",
    "def train(epoch):\n",
    "    tr_loss, tr_accuracy = 0, 0\n",
    "    nb_tr_examples, nb_tr_steps = 0, 0\n",
    "    tr_preds, tr_labels = [], []\n",
    "    # put model in training mode\n",
    "    model.train()\n",
    "    \n",
    "    for idx, batch in enumerate(training_loader):\n",
    "        \n",
    "        ids = batch['ids'].to(device, dtype = torch.long)\n",
    "        mask = batch['mask'].to(device, dtype = torch.long)\n",
    "        targets = batch['targets'].to(device, dtype = torch.long)\n",
    "\n",
    "        loss, tr_logits = model(input_ids=ids, attention_mask=mask, labels=targets)\n",
    "        tr_loss += loss.item()\n",
    "\n",
    "        nb_tr_steps += 1\n",
    "        nb_tr_examples += targets.size(0)\n",
    "        \n",
    "        if idx % 100==0:\n",
    "            loss_step = tr_loss/nb_tr_steps\n",
    "            print(f\"Training loss per 100 training steps: {loss_step}\")\n",
    "           \n",
    "        # compute training accuracy\n",
    "        flattened_targets = targets.view(-1) # shape (batch_size * seq_len,)\n",
    "        active_logits = tr_logits.view(-1, model.num_labels) # shape (batch_size * seq_len, num_labels)\n",
    "        flattened_predictions = torch.argmax(active_logits, axis=1) # shape (batch_size * seq_len,)\n",
    "        # now, use mask to determine where we should compare predictions with targets (includes [CLS] and [SEP] token predictions)\n",
    "        active_accuracy = mask.view(-1) == 1 # active accuracy is also of shape (batch_size * seq_len,)\n",
    "        targets = torch.masked_select(flattened_targets, active_accuracy)\n",
    "        predictions = torch.masked_select(flattened_predictions, active_accuracy)\n",
    "        \n",
    "        tr_preds.extend(predictions)\n",
    "        tr_labels.extend(targets)\n",
    "        \n",
    "        tmp_tr_accuracy = accuracy_score(targets.cpu().numpy(), predictions.cpu().numpy())\n",
    "        tr_accuracy += tmp_tr_accuracy\n",
    "    \n",
    "        # gradient clipping\n",
    "        torch.nn.utils.clip_grad_norm_(\n",
    "            parameters=model.parameters(), max_norm=MAX_GRAD_NORM\n",
    "        )\n",
    "        \n",
    "        # backward pass\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "    epoch_loss = tr_loss / nb_tr_steps\n",
    "    tr_accuracy = tr_accuracy / nb_tr_steps\n",
    "    print(f\"Training loss epoch: {epoch_loss}\")\n",
    "    print(f\"Training accuracy epoch: {tr_accuracy}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "k2dsCyP7dcF3"
   },
   "source": [
    "And let's train the model!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 521
    },
    "colab_type": "code",
    "id": "y07Ybw8rZeZ7",
    "outputId": "25bec966-fa2c-461b-a1cf-647fb26b143f",
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training epoch: 1\n",
      "Training loss per 100 training steps: 2.489778757095337\n",
      "Training loss per 100 training steps: 0.8853663772639662\n",
      "Training loss per 100 training steps: 0.6615404157125535\n",
      "Training loss per 100 training steps: 0.554439210199884\n",
      "Training loss per 100 training steps: 0.4880749371478757\n",
      "Training loss per 100 training steps: 0.4382783724050857\n",
      "Training loss per 100 training steps: 0.40183136461553676\n",
      "Training loss per 100 training steps: 0.3772019159776848\n",
      "Training loss per 100 training steps: 0.3559265573661351\n",
      "Training loss per 100 training steps: 0.3382058288451371\n",
      "Training loss per 100 training steps: 0.32617049882156807\n",
      "Training loss per 100 training steps: 0.31556622421099967\n",
      "Training loss per 100 training steps: 0.3042621851730456\n",
      "Training loss per 100 training steps: 0.2951039669553475\n",
      "Training loss per 100 training steps: 0.2885802547057837\n",
      "Training loss per 100 training steps: 0.28103069619720505\n",
      "Training loss per 100 training steps: 0.27402849008313507\n",
      "Training loss per 100 training steps: 0.26728194428614394\n",
      "Training loss per 100 training steps: 0.2621294350721674\n",
      "Training loss per 100 training steps: 0.25749910095983874\n",
      "Training loss per 100 training steps: 0.25315395947908803\n",
      "Training loss per 100 training steps: 0.25034797743321\n",
      "Training loss per 100 training steps: 0.24694465017872866\n",
      "Training loss per 100 training steps: 0.24384244190677126\n",
      "Training loss per 100 training steps: 0.24054359724789373\n",
      "Training loss per 100 training steps: 0.2380536368950123\n",
      "Training loss per 100 training steps: 0.23545476084650965\n",
      "Training loss per 100 training steps: 0.23244921516014008\n",
      "Training loss per 100 training steps: 0.22921396398275787\n",
      "Training loss per 100 training steps: 0.2258582871677903\n",
      "Training loss per 100 training steps: 0.2237785360561755\n",
      "Training loss per 100 training steps: 0.22219097535166457\n",
      "Training loss per 100 training steps: 0.21937036143047914\n",
      "Training loss per 100 training steps: 0.21693665391230327\n",
      "Training loss per 100 training steps: 0.21457680911622598\n",
      "Training loss per 100 training steps: 0.21255231931777946\n",
      "Training loss per 100 training steps: 0.21045138725793153\n",
      "Training loss per 100 training steps: 0.2090586207417524\n",
      "Training loss per 100 training steps: 0.20748627698673752\n",
      "Training loss per 100 training steps: 0.20592153333394886\n",
      "Training loss per 100 training steps: 0.2042281029585789\n",
      "Training loss per 100 training steps: 0.2025867179522127\n",
      "Training loss per 100 training steps: 0.20164298638524086\n",
      "Training loss per 100 training steps: 0.20077074291797342\n",
      "Training loss per 100 training steps: 0.19946115016065397\n",
      "Training loss per 100 training steps: 0.1984474846129106\n",
      "Training loss per 100 training steps: 0.19716404113130936\n",
      "Training loss per 100 training steps: 0.19602193017973601\n",
      "Training loss per 100 training steps: 0.19467369015117142\n",
      "Training loss per 100 training steps: 0.19362043105948498\n",
      "Training loss per 100 training steps: 0.19259222732191852\n",
      "Training loss per 100 training steps: 0.1912094354522253\n",
      "Training loss per 100 training steps: 0.18998318648772838\n",
      "Training loss per 100 training steps: 0.18897544341890074\n",
      "Training loss per 100 training steps: 0.18812571352200644\n",
      "Training loss per 100 training steps: 0.1871206830665818\n",
      "Training loss per 100 training steps: 0.18636027630547405\n",
      "Training loss per 100 training steps: 0.18553527820368726\n",
      "Training loss per 100 training steps: 0.18471234588403146\n",
      "Training loss per 100 training steps: 0.18439051741959048\n",
      "Training loss per 100 training steps: 0.18370074144464624\n",
      "Training loss per 100 training steps: 0.18307926314895145\n",
      "Training loss per 100 training steps: 0.18206197839365318\n",
      "Training loss per 100 training steps: 0.1814760097080862\n",
      "Training loss per 100 training steps: 0.1807653515840392\n",
      "Training loss per 100 training steps: 0.17984701719349505\n",
      "Training loss per 100 training steps: 0.17911162204560616\n",
      "Training loss per 100 training steps: 0.17864168832182675\n",
      "Training loss per 100 training steps: 0.17796534106626022\n",
      "Training loss per 100 training steps: 0.17737074654391868\n",
      "Training loss per 100 training steps: 0.17666300702465618\n",
      "Training loss per 100 training steps: 0.17595806113670379\n",
      "Training loss per 100 training steps: 0.17545315528951572\n",
      "Training loss per 100 training steps: 0.17499334982272588\n",
      "Training loss per 100 training steps: 0.17439171155924604\n",
      "Training loss per 100 training steps: 0.17374772254908138\n",
      "Training loss per 100 training steps: 0.17333722603446078\n",
      "Training loss per 100 training steps: 0.1726307541946412\n",
      "Training loss per 100 training steps: 0.17214256509510667\n",
      "Training loss per 100 training steps: 0.17137341667111072\n",
      "Training loss per 100 training steps: 0.1708224866951997\n",
      "Training loss per 100 training steps: 0.1702852733746236\n",
      "Training loss per 100 training steps: 0.1698461759778076\n",
      "Training loss per 100 training steps: 0.16935828631911967\n",
      "Training loss per 100 training steps: 0.1689116654104395\n",
      "Training loss per 100 training steps: 0.16865811368572456\n",
      "Training loss per 100 training steps: 0.1680868287597195\n",
      "Training loss per 100 training steps: 0.16769062422443712\n",
      "Training loss per 100 training steps: 0.16742355229466166\n",
      "Training loss per 100 training steps: 0.16707693207979707\n",
      "Training loss per 100 training steps: 0.1665783482122599\n",
      "Training loss per 100 training steps: 0.16630573462046713\n",
      "Training loss per 100 training steps: 0.1658349263041819\n",
      "Training loss per 100 training steps: 0.16518718661170323\n",
      "Training loss per 100 training steps: 0.16461824608948677\n",
      "Training loss per 100 training steps: 0.1641890350354777\n",
      "Training loss epoch: 0.1641632655221446\n",
      "Training accuracy epoch: 0.951427250396801\n"
     ]
    }
   ],
   "source": [
    "for epoch in range(EPOCHS):\n",
    "    print(f\"Training epoch: {epoch + 1}\")\n",
    "    train(epoch)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "r4jcSOJr680a"
   },
   "source": [
    "#### **Evaluating the model**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "rYUTuOEUdfFJ"
   },
   "source": [
    "Now that we've trained our model, we can evaluate its performance on the held-out test set (which is 20% of the data). Note that here, no gradient updates are performed, the model just outputs its logits. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "RIVVfFHi7Aw7"
   },
   "outputs": [],
   "source": [
    "def valid(model, testing_loader):\n",
    "    # put model in evaluation mode\n",
    "    model.eval()\n",
    "    \n",
    "    eval_loss, eval_accuracy = 0, 0\n",
    "    nb_eval_examples, nb_eval_steps = 0, 0\n",
    "    eval_preds, eval_labels = [], []\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        for idx, batch in enumerate(testing_loader):\n",
    "            \n",
    "            ids = batch['ids'].to(device, dtype = torch.long)\n",
    "            mask = batch['mask'].to(device, dtype = torch.long)\n",
    "            targets = batch['targets'].to(device, dtype = torch.long)\n",
    "            \n",
    "            loss, eval_logits = model(input_ids=ids, attention_mask=mask, labels=targets)\n",
    "            \n",
    "            eval_loss += loss.item()\n",
    "\n",
    "            nb_eval_steps += 1\n",
    "            nb_eval_examples += targets.size(0)\n",
    "        \n",
    "            if idx % 100==0:\n",
    "                loss_step = eval_loss/nb_eval_steps\n",
    "                print(f\"Validation loss per 100 evaluation steps: {loss_step}\")\n",
    "              \n",
    "            # compute evaluation accuracy\n",
    "            flattened_targets = targets.view(-1) # shape (batch_size * seq_len,)\n",
    "            active_logits = eval_logits.view(-1, model.num_labels) # shape (batch_size * seq_len, num_labels)\n",
    "            flattened_predictions = torch.argmax(active_logits, axis=1) # shape (batch_size * seq_len,)\n",
    "            # now, use mask to determine where we should compare predictions with targets (includes [CLS] and [SEP] token predictions)\n",
    "            active_accuracy = mask.view(-1) == 1 # active accuracy is also of shape (batch_size * seq_len,)\n",
    "            targets = torch.masked_select(flattened_targets, active_accuracy)\n",
    "            predictions = torch.masked_select(flattened_predictions, active_accuracy)\n",
    "            \n",
    "            eval_labels.extend(targets)\n",
    "            eval_preds.extend(predictions)\n",
    "            \n",
    "            tmp_eval_accuracy = accuracy_score(targets.cpu().numpy(), predictions.cpu().numpy())\n",
    "            eval_accuracy += tmp_eval_accuracy\n",
    "    \n",
    "    #print(eval_labels)\n",
    "    #print(eval_preds)\n",
    "\n",
    "    labels = [ids_to_labels[id.item()] for id in eval_labels]\n",
    "    predictions = [ids_to_labels[id.item()] for id in eval_preds]\n",
    "\n",
    "    #print(labels)\n",
    "    #print(predictions)\n",
    "    \n",
    "    eval_loss = eval_loss / nb_eval_steps\n",
    "    eval_accuracy = eval_accuracy / nb_eval_steps\n",
    "    print(f\"Validation Loss: {eval_loss}\")\n",
    "    print(f\"Validation Accuracy: {eval_accuracy}\")\n",
    "\n",
    "    return labels, predictions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "zJaONluRdq-e"
   },
   "source": [
    "As we can see below, performance is quite good! Accuracy on the test test is > 93%."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "2BrxRjvxApY8",
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation loss per 100 evaluation steps: 0.16408051550388336\n",
      "Validation loss per 100 evaluation steps: 0.14635943754179642\n",
      "Validation loss per 100 evaluation steps: 0.12696247273796604\n",
      "Validation loss per 100 evaluation steps: 0.13026319034637396\n",
      "Validation loss per 100 evaluation steps: 0.12882538416952433\n",
      "Validation loss per 100 evaluation steps: 0.12492606570275441\n",
      "Validation loss per 100 evaluation steps: 0.12379682637610087\n",
      "Validation loss per 100 evaluation steps: 0.12447704667265005\n",
      "Validation loss per 100 evaluation steps: 0.12321146057993371\n",
      "Validation loss per 100 evaluation steps: 0.12356419575621078\n",
      "Validation loss per 100 evaluation steps: 0.12217774208278984\n",
      "Validation loss per 100 evaluation steps: 0.12188346979894411\n",
      "Validation loss per 100 evaluation steps: 0.12174783734095629\n",
      "Validation loss per 100 evaluation steps: 0.12099008415253432\n",
      "Validation loss per 100 evaluation steps: 0.12045939129233067\n",
      "Validation loss per 100 evaluation steps: 0.12171252334249522\n",
      "Validation loss per 100 evaluation steps: 0.12189806149098519\n",
      "Validation loss per 100 evaluation steps: 0.12161771707661911\n",
      "Validation loss per 100 evaluation steps: 0.12101843894540365\n",
      "Validation loss per 100 evaluation steps: 0.12079252161901063\n",
      "Validation loss per 100 evaluation steps: 0.11974893306831547\n",
      "Validation loss per 100 evaluation steps: 0.12206078356969759\n",
      "Validation loss per 100 evaluation steps: 0.12306439843317733\n",
      "Validation loss per 100 evaluation steps: 0.12233479616827396\n",
      "Validation loss per 100 evaluation steps: 0.1215667084076492\n",
      "Validation loss per 100 evaluation steps: 0.12082167003658086\n",
      "Validation loss per 100 evaluation steps: 0.12017407765367091\n",
      "Validation loss per 100 evaluation steps: 0.11968497871769102\n",
      "Validation loss per 100 evaluation steps: 0.11945383894316348\n",
      "Validation loss per 100 evaluation steps: 0.11997537971993304\n",
      "Validation loss per 100 evaluation steps: 0.12008571720221466\n",
      "Validation loss per 100 evaluation steps: 0.12000624452732168\n",
      "Validation loss per 100 evaluation steps: 0.12104994785082804\n",
      "Validation loss per 100 evaluation steps: 0.12158328771310076\n",
      "Validation loss per 100 evaluation steps: 0.12073075326625086\n",
      "Validation loss per 100 evaluation steps: 0.12104582718923787\n",
      "Validation loss per 100 evaluation steps: 0.12103564786278675\n",
      "Validation loss per 100 evaluation steps: 0.12111494211068337\n",
      "Validation loss per 100 evaluation steps: 0.1210244209080354\n",
      "Validation loss per 100 evaluation steps: 0.12098946495988859\n",
      "Validation loss per 100 evaluation steps: 0.12125087661863285\n",
      "Validation loss per 100 evaluation steps: 0.12137658662944557\n",
      "Validation loss per 100 evaluation steps: 0.1208742182077268\n",
      "Validation loss per 100 evaluation steps: 0.12089782744794725\n",
      "Validation loss per 100 evaluation steps: 0.12069129760921386\n",
      "Validation loss per 100 evaluation steps: 0.12065088340305312\n",
      "Validation loss per 100 evaluation steps: 0.12039781877801838\n",
      "Validation loss per 100 evaluation steps: 0.12000540961992398\n",
      "Validation Loss: 0.1200615793137732\n",
      "Validation Accuracy: 0.9625447726461819\n"
     ]
    }
   ],
   "source": [
    "labels, predictions = valid(model, testing_loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "SAznLDwx_U2X"
   },
   "source": [
    "However, the accuracy metric is misleading, as a lot of labels are \"outside\" (O), even after omitting predictions on the [PAD] tokens. What is important is looking at the precision, recall and f1-score of the individual tags. For this, we use the seqeval Python library: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "0jDNXrjr-6BW",
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "           precision    recall  f1-score   support\n",
      "\n",
      "      org       0.72      0.58      0.64      6531\n",
      "      per       0.75      0.79      0.77      5196\n",
      "      geo       0.79      0.92      0.85     11232\n",
      "      gpe       0.94      0.92      0.93      3293\n",
      "      tim       0.87      0.81      0.84      4360\n",
      "\n",
      "micro avg       0.80      0.81      0.80     30612\n",
      "macro avg       0.80      0.81      0.80     30612\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from seqeval.metrics import classification_report\n",
    "\n",
    "print(classification_report(labels, predictions))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "4Gz-wHAw3xMk"
   },
   "source": [
    "#### **Inference**\n",
    "\n",
    "The fun part is when we can quickly test the model on new, unseen sentences. \n",
    "Here, we use the prediction of the **first word piece of every word**. Note that the function we used to train our model (`tokenze_and_preserve_labels`) propagated the label to all subsequent word pieces (so you could for example also perform a majority vote on the predicted labels of all word pieces of a word).\n",
    "\n",
    "*In other words, the code below does not take into account when predictions of different word pieces that belong to the same word do not match.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "zPDla1mmZiax",
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "india has a capital called mumbai . on wednesday , the president will give a presentation\n",
      "['B-geo', 'O', 'O', 'O', 'O', 'B-geo', 'O', 'O', 'B-tim', 'O', 'O', 'O', 'O', 'O', 'O', 'O']\n"
     ]
    }
   ],
   "source": [
    "sentence = \"India has a capital called Mumbai. On wednesday, the president will give a presentation\"\n",
    "\n",
    "inputs = tokenizer(sentence, padding='max_length', truncation=True, max_length=MAX_LEN, return_tensors=\"pt\")\n",
    "\n",
    "# move to gpu\n",
    "ids = inputs[\"input_ids\"].to(device)\n",
    "mask = inputs[\"attention_mask\"].to(device)\n",
    "# forward pass\n",
    "outputs = model(ids, mask)\n",
    "logits = outputs[0]\n",
    "\n",
    "active_logits = logits.view(-1, model.num_labels) # shape (batch_size * seq_len, num_labels)\n",
    "flattened_predictions = torch.argmax(active_logits, axis=1) # shape (batch_size*seq_len,) - predictions at the token level\n",
    "\n",
    "tokens = tokenizer.convert_ids_to_tokens(ids.squeeze().tolist())\n",
    "token_predictions = [ids_to_labels[i] for i in flattened_predictions.cpu().numpy()]\n",
    "wp_preds = list(zip(tokens, token_predictions)) # list of tuples. Each tuple = (wordpiece, prediction)\n",
    "\n",
    "word_level_predictions = []\n",
    "for pair in wp_preds:\n",
    "  if (pair[0].startswith(\" ##\")) or (pair[0] in ['[CLS]', '[SEP]', '[PAD]']):\n",
    "    # skip prediction\n",
    "    continue\n",
    "  else:\n",
    "    word_level_predictions.append(pair[1])\n",
    "\n",
    "# we join tokens, if they are not special ones\n",
    "str_rep = \" \".join([t[0] for t in wp_preds if t[0] not in ['[CLS]', '[SEP]', '[PAD]']]).replace(\" ##\", \"\")\n",
    "print(str_rep)\n",
    "print(word_level_predictions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "sqDklprSqB5d"
   },
   "source": [
    "#### **Saving the model for future use**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "VuUdX_fImswO"
   },
   "source": [
    "Finally, let's save the vocabulary, model weights and the model's configuration file to a directory, so that they can be re-loaded using the `from_pretrained()` class method.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "sDZtSsKKntuI",
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "All files saved\n",
      "This tutorial is completed\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "\n",
    "directory = \"./model\"\n",
    "\n",
    "if not os.path.exists(directory):\n",
    "    os.makedirs(directory)\n",
    "\n",
    "# save vocabulary of the tokenizer\n",
    "tokenizer.save_vocabulary(directory)\n",
    "# save the model weights and its configuration file\n",
    "model.save_pretrained(directory)\n",
    "print('All files saved')\n",
    "print('This tutorial is completed')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Legacy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def prepare_sentence(sentence, tokenizer, maxlen):    \n",
    "      # step 1: tokenize the sentence\n",
    "      tokenized_sentence = tokenizer.tokenize(sentence)\n",
    "      \n",
    "      # step 2: add special tokens \n",
    "      tokenized_sentence = [\"[CLS]\"] + tokenized_sentence + [\"[SEP]\"] \n",
    "\n",
    "      # step 3: truncating/padding\n",
    "      if (len(tokenized_sentence) > maxlen):\n",
    "        # truncate\n",
    "        tokenized_sentence = tokenized_sentence[:maxlen]\n",
    "      else:\n",
    "        # pad\n",
    "        tokenized_sentence = tokenized_sentence + ['[PAD]'for _ in range(maxlen - len(tokenized_sentence))]\n",
    "\n",
    "      # step 4: obtain the attention mask\n",
    "      attn_mask = [1 if tok != '[PAD]' else 0 for tok in tokenized_sentence]\n",
    "      \n",
    "      # step 5: convert tokens to input ids\n",
    "      ids = tokenizer.convert_tokens_to_ids(tokenized_sentence)\n",
    "      \n",
    "      return {\n",
    "            'ids': torch.tensor(ids, dtype=torch.long),\n",
    "            'mask': torch.tensor(attn_mask, dtype=torch.long),\n",
    "            #'token_type_ids': torch.tensor(token_ids, dtype=torch.long),\n",
    "      }"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "authorship_tag": "ABX9TyMxMQ3bEovJTibndikJm1d9",
   "collapsed_sections": [],
   "include_colab_link": true,
   "name": "Custom Named Entity Recognition with BERT.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
